from torchvision.datasets import FashionMNIST
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import torch


def load_dataset(batch_size=32):
    mnist_test = FashionMNIST(root='~/Datasets/FashionMNIST',
                       train=False, download=True,
                       transform=transforms.ToTensor())
    mnist_train = FashionMNIST(root='~/Datasets/FashionMNIST',
                        train=True, download=True,
                        transform=transforms.ToTensor())
    train_iter = DataLoader(mnist_train, batch_size=batch_size, shuffle=True)
    test_iter = DataLoader(mnist_test, batch_size=batch_size, shuffle=False)
    return train_iter, test_iter


def evaluate(data_iter, model):
    model.eval()
    with torch.no_grad():
        acc_sum, n = 0.0, 0
        for x, y in data_iter:
            logits = model(x)
            acc_sum += (logits.argmax(1) == y).float().sum().item()
            n += len(y)
        model.train()
        return acc_sum / n